package handler

import (
	"fmt"
	"gitee.com/simple-set/simple.io/src/event"
	"gitee.com/simple-set/simple.io/src/process/codec"
	"gitee.com/simple-set/simple.io/src/process/socks"
	"gitee.com/simple-set/simple.io/src/socket"
	"log/slog"
	"net"
)

type SocksAssociateHandler struct {
	client *Client
}

func NewSocksAssociateHandler() *SocksAssociateHandler {
	return &SocksAssociateHandler{client: NewClient()}
}

func (s *SocksAssociateHandler) Disconnect(session *event.Session) error {
	if udpServer := session.GetAttr("udpServer"); udpServer != nil {
		if server, ok := udpServer.(socket.Server); ok {
			_ = server.Close()
		}
	}
	return nil
}

func (s *SocksAssociateHandler) Associate(request *socks.Request, response *socks.Response) error {
	udpServer, err := s.createUdpServer(fmt.Sprintf("%s:", s.getListenAddr(request)), response)
	if err != nil {
		return err
	}
	if err := response.SetAddrAndPort(udpServer.ListenAddr()); err != nil {
		_ = response.Reply(socks.ServerFailure)
		return err
	}
	if err := response.Reply(socks.SuccessReply); err != nil {
		return err
	}

	request.Session().SetAttr("udpServer", udpServer)
	go s.relay(udpServer, request)
	return nil
}

func (s *SocksAssociateHandler) relay(udpServer socket.Server, request *socks.Request) {
	for {
		sock, err := udpServer.Listen()
		if err != nil {
			_ = udpServer.Close()
			slog.Error("sock.Associate UdpServer listening exception", "err", err)
			return
		}
		if !s.udpRelayIsAllow(sock.RemoteAddr().(*net.UDPAddr), request) {
			slog.Warn("udpRelay Address mismatch", "dstAddr", sock.RemoteAddr().String(), "expectation", request.GetDstAddrString())
			_ = sock.Close()
			continue
		}
		go s.doRelay(sock)
	}
}

func (s *SocksAssociateHandler) doRelay(sock *socket.Socket) {
	defer func() {
		_ = sock.Close()
	}()

	udpRelayRequest := socks.NewUdpRelay(codec.NewReadByteBuf(sock))
	if err := udpRelayRequest.Decode(); err != nil {
		slog.Error("udpRelayRequest data decode failed", "err", err)
		return
	}
	udpRelayPayload, err := udpRelayRequest.Payload()
	if err != nil {
		slog.Error("udpRelayRequest get Payload failed", "err", err)
		return
	}
	// 转发udp中继的有效载荷到目标服务器
	receivedData, err := s.udpRelaTransfer(udpRelayRequest.GetDstAddrString(), udpRelayPayload)
	if err != nil {
		slog.Error("udpRelaTransfer failed", "err", err)
		return
	}

	slog.Debug(fmt.Sprintf("Socks.Associate %s -> %s", sock.RemoteAddr().String(), udpRelayRequest.GetDstAddr().String()))
	udpRelayResponse := socks.NewUdpRelay(codec.NewWriteByteBuf(sock))
	udpRelayResponse.SetPayload(receivedData)
	udpRelayResponse.SetAddr(udpRelayRequest.GetDstAddr())
	if err := udpRelayResponse.Encode(); err != nil {
		slog.Error("udpRelayRequest encode failed", "err", err)
	}
}

func (s *SocksAssociateHandler) createUdpServer(addr string, response *socks.Response) (socket.Server, error) {
	udpServer := socket.NewUdpServer(addr)
	if err := udpServer.Bind(); err != nil {
		_ = response.Reply(socks.NetworkUnreachable)
		return nil, err
	}
	return udpServer, nil
}

func (s *SocksAssociateHandler) getListenAddr(request *socks.Request) string {
	if tcpAddr, ok := request.Session().Socket.LocalAddr().(*net.TCPAddr); ok {
		return tcpAddr.IP.String()
	}
	return request.Session().Socket.LocalAddr().(*net.UDPAddr).IP.String()
}

func (s *SocksAssociateHandler) udpRelayIsAllow(relayAddr *net.UDPAddr, request *socks.Request) bool {
	if net.IP(request.DstAddr).IsUnspecified() {
		return relayAddr.IP.Equal(request.Session().Socket.RemoteAddr().(*net.TCPAddr).IP)
	} else {
		return relayAddr.IP.Equal(request.DstAddr)
	}
}

func (s *SocksAssociateHandler) udpRelaTransfer(addr string, sendData []byte) ([]byte, error) {
	udpClientSocket, err := s.client.DialUdpSocket(addr)
	if err != nil {
		return nil, err
	}
	defer func() {
		_ = udpClientSocket.Close()
	}()
	if err := udpClientSocket.WriterBytes(sendData); err != nil {
		return nil, err
	}

	// 从目标服务器读取回复, 受限MTU大小, UDP接收缓冲区1500字节
	receivedData := make([]byte, 1500)
	if n, err := udpClientSocket.Read(receivedData); err != nil {
		return nil, err
	} else {
		return receivedData[:n], nil
	}
}
